use async_trait::async_trait;
use fluviox_shared::{Result, FluvioXError, models::DataRecord};
use crate::traits::{SmartModule, Transform, SmartModuleType, ProcessResult};
use serde_json::Value;
use jsonpath_lib as jsonpath;
use tracing::{info, error, warn};

/// JSON转换器
/// 使用JSONPath进行数据提取和转换
pub struct JsonTransformer {
    name: String,
    config: Option<JsonTransformerConfig>,
}

/// JSON转换器配置
#[derive(Debug, Clone)]
pub struct JsonTransformerConfig {
    /// JSONPath提取规则
    pub extractions: Vec<ExtractionRule>,
    /// 是否展平嵌套对象
    pub flatten: bool,
    /// 展平分隔符
    pub flatten_separator: String,
    /// 是否移除空值
    pub remove_null: bool,
    /// 输出模式
    pub output_mode: OutputMode,
}

/// 提取规则
#[derive(Debug, Clone)]
pub struct ExtractionRule {
    /// 源JSONPath
    pub source_path: String,
    /// 目标字段名
    pub target_field: String,
    /// 是否必需
    pub required: bool,
    /// 默认值（如果提取失败）
    pub default_value: Option<Value>,
}

/// 输出模式
#[derive(Debug, Clone)]
pub enum OutputMode {
    /// 保留原始数据，添加提取的字段
    Merge,
    /// 只输出提取的字段
    ExtractOnly,
    /// 替换原始数据
    Replace,
}

impl JsonTransformer {
    pub fn new() -> Self {
        Self {
            name: "json-transformer".to_string(),
            config: None,
        }
    }
    
    /// 展平JSON对象
    fn flatten_object(&self, obj: &serde_json::Map<String, Value>, separator: &str) -> serde_json::Map<String, Value> {
        let mut result = serde_json::Map::new();
        self.flatten_recursive(obj, "", separator, &mut result);
        result
    }
    
    /// 递归展平
    fn flatten_recursive(
        &self,
        obj: &serde_json::Map<String, Value>,
        prefix: &str,
        separator: &str,
        result: &mut serde_json::Map<String, Value>,
    ) {
        for (key, value) in obj {
            let new_key = if prefix.is_empty() {
                key.clone()
            } else {
                format!("{}{}{}", prefix, separator, key)
            };
            
            match value {
                Value::Object(nested_obj) => {
                    self.flatten_recursive(nested_obj, &new_key, separator, result);
                },
                _ => {
                    result.insert(new_key, value.clone());
                }
            }
        }
    }
    
    /// 应用JSONPath提取
    fn apply_extractions(&self, data: &Value, config: &JsonTransformerConfig) -> Result<serde_json::Map<String, Value>> {
        let mut extracted = serde_json::Map::new();
        
        for rule in &config.extractions {
            match jsonpath::select(data, &rule.source_path) {
                Ok(results) => {
                    if results.is_empty() {
                        if rule.required {
                            return Err(FluvioXError::Internal(
                                format!("必需字段 '{}' 提取失败", rule.target_field)
                            ));
                        } else if let Some(default) = &rule.default_value {
                            extracted.insert(rule.target_field.clone(), default.clone());
                        }
                    } else {
                        // 如果有多个结果，取第一个
                        let value = if results.len() == 1 {
                            results[0].clone()
                        } else {
                            Value::Array(results.into_iter().cloned().collect())
                        };
                        extracted.insert(rule.target_field.clone(), value);
                    }
                },
                Err(e) => {
                    error!("JSONPath提取失败: {} - {}", rule.source_path, e);
                    if rule.required {
                        return Err(FluvioXError::Internal(
                            format!("JSONPath提取失败: {}", e)
                        ));
                    } else if let Some(default) = &rule.default_value {
                        extracted.insert(rule.target_field.clone(), default.clone());
                    }
                }
            }
        }
        
        Ok(extracted)
    }
    
    /// 移除空值
    fn remove_null_values(&self, mut obj: serde_json::Map<String, Value>) -> serde_json::Map<String, Value> {
        obj.retain(|_, v| !v.is_null());
        obj
    }
    
    /// 应用转换
    fn apply_transform(&self, data: Value, config: &JsonTransformerConfig) -> Result<Value> {
        let mut result_obj = serde_json::Map::new();
        
        // 应用JSONPath提取
        let extracted = self.apply_extractions(&data, config)?;
        
        match config.output_mode {
            OutputMode::ExtractOnly => {
                result_obj = extracted;
            },
            OutputMode::Replace => {
                result_obj = extracted;
            },
            OutputMode::Merge => {
                // 先添加原始数据
                if let Value::Object(original_obj) = data {
                    result_obj = original_obj;
                }
                // 然后添加提取的数据
                for (key, value) in extracted {
                    result_obj.insert(key, value);
                }
            },
        }
        
        // 展平处理
        if config.flatten {
            result_obj = self.flatten_object(&result_obj, &config.flatten_separator);
        }
        
        // 移除空值
        if config.remove_null {
            result_obj = self.remove_null_values(result_obj);
        }
        
        Ok(Value::Object(result_obj))
    }
}

#[async_trait]
impl SmartModule for JsonTransformer {
    fn name(&self) -> &str {
        &self.name
    }
    
    fn version(&self) -> &str {
        "1.0.0"
    }
    
    fn description(&self) -> &str {
        "JSON转换器，支持JSONPath提取、对象展平和数据重构"
    }
    
    fn module_type(&self) -> SmartModuleType {
        SmartModuleType::Transform
    }
    
    async fn initialize(&mut self, config: Value) -> Result<()> {
        info!("初始化JSON转换器");
        
        // 解析提取规则
        let extractions = config.get("extractions")
            .and_then(|v| v.as_array())
            .map(|arr| {
                arr.iter()
                    .filter_map(|item| {
                        if let Value::Object(rule_obj) = item {
                            let source_path = rule_obj.get("source_path")?.as_str()?.to_string();
                            let target_field = rule_obj.get("target_field")?.as_str()?.to_string();
                            let required = rule_obj.get("required")
                                .and_then(|v| v.as_bool())
                                .unwrap_or(false);
                            let default_value = rule_obj.get("default_value").cloned();
                            
                            Some(ExtractionRule {
                                source_path,
                                target_field,
                                required,
                                default_value,
                            })
                        } else {
                            None
                        }
                    })
                    .collect()
            })
            .unwrap_or_default();
        
        let flatten = config.get("flatten")
            .and_then(|v| v.as_bool())
            .unwrap_or(false);
        
        let flatten_separator = config.get("flatten_separator")
            .and_then(|v| v.as_str())
            .unwrap_or(".")
            .to_string();
        
        let remove_null = config.get("remove_null")
            .and_then(|v| v.as_bool())
            .unwrap_or(false);
        
        let output_mode = config.get("output_mode")
            .and_then(|v| v.as_str())
            .map(|s| match s {
                "extract_only" => OutputMode::ExtractOnly,
                "replace" => OutputMode::Replace,
                _ => OutputMode::Merge,
            })
            .unwrap_or(OutputMode::Merge);
        
        self.config = Some(JsonTransformerConfig {
            extractions,
            flatten,
            flatten_separator,
            remove_null,
            output_mode,
        });
        
        info!("JSON转换器初始化完成");
        Ok(())
    }
    
    async fn process(&mut self, mut record: DataRecord) -> Result<Vec<ProcessResult>> {
        let config = self.config.as_ref()
            .ok_or_else(|| FluvioXError::Internal("转换器未初始化".to_string()))?;
        
        // 应用JSON转换
        record.data = self.apply_transform(record.data, config)?;
        
        Ok(vec![ProcessResult::Transform(record)])
    }
}

#[async_trait]
impl Transform for JsonTransformer {
    async fn transform(&self, mut record: DataRecord) -> Result<DataRecord> {
        let config = self.config.as_ref()
            .ok_or_else(|| FluvioXError::Internal("转换器未初始化".to_string()))?;
        
        record.data = self.apply_transform(record.data, config)?;
        Ok(record)
    }
}

impl Default for JsonTransformer {
    fn default() -> Self {
        Self::new()
    }
}
